import argparse


class Config(object):
    def __init__(self):
        self.project_name = "Image Recover"
        self.author_name = "Luotianhang"

        self.use_cuda = True
        self.seed = 123
        # for Ubuntu
        # self.train_data_path = "/home/ImageNet/train/"
        # self.test_data_path = "/home/ImageNet/val/"
        # for Windows
        self.train_data_path='G:/datasets/VOC/VOC2007/JPEGImages'
        self.test_data_path='G:/datasets/VOC/VOC2012/JPEGImages'
        self.resume = True
        self.pretrained_model = "./weights/best.pth"

        self.lr = 0.001
        self.epochs = 200
        self.train_batch_size = 18
        self.test_batch_size = 18
        self.weight_decay = 1e-4
        self.momentum = 0.5
        self.save = True
        self.image_size = 256


def parser_config():
    parser = argparse.ArgumentParser(description="PyTorch")
    #####################################################################################
    # TODO project information
    parser.add_argument("--project_name", type=str, default="Image Recover")
    parser.add_argument("--author_name", type=str, default="Luotianhang")
    #####################################################################################
    # TODO train.py
    # TODO basic operation for the model training
    parser.add_argument('--lr', type=float, default=0.001)
    parser.add_argument('--epochs', type=int, default=200)
    parser.add_argument('--train_batch_size', type=int, default=8)
    parser.add_argument('--test_batch_size', type=int, default=8)
    parser.add_argument('--weight_decay', type=float, default=1e-4)
    parser.add_argument('--momentum', type=float, default=0.5)
    parser.add_argument('--save', type=bool, default=True)
    parser.add_argument('--resume', type=bool, default=True)
    parser.add_argument('--use_cuda', type=bool, default=True)
    parser.add_argument('--seed', type=int, default=123)
    parser.add_argument('--pretrained_model', type=str, default="./weights/best.pth")
    #####################################################################################
    # TODO datalist.py
    parser.add_argument("--image_size", type=int, default=256)
    # for ubuntu to use
    parser.add_argument("--train_data_path", type=str, default="G:/datasets/VOC/VOC2007/JPEGImages")
    parser.add_argument("--test_data_path", type=str, default="G:/datasets/VOC/VOC2012/JPEGImages")
    # for windows to use
    # parser.add_argument("--train_data_path", type=str, default="./")
    # parser.add_argument("--test_data_path", type=str, default="./")
    #####################################################################################
    # TODO DDP model
    parser.add_argument("--device", default='', help="cuda device")
    parser.add_argument("--local_rank", type=int, default=-1, help='DDP parameter,do not modify')

    return parser.parse_args()
